#pragma once
#include<opencv2/highgui/highgui.hpp>
#include<opencv2/imgproc/imgproc.hpp>
#include<opencv2/dnn.hpp>
#include<opencv2/calib3d.hpp>

#include<iostream>
#include<fstream>
#include<vector>

#define INFO_OK							0     // 正常
#define INFO_LOAD_NET_ERR				1     // 加载模型失败
#define INFO_LOAD_CLASS_ERR				2     // 加载类别失败

struct YoloDetection
{
	int class_id;
	float confidence;
	cv::Rect box;
};

struct YoloParam
{
	float score_threshold;
	float nms_threshold;
	float confidence_threshold;
};


class YoloDet
{
public:
	YoloDet();
	~YoloDet();
	int init(std::string model_path, std::string class_path, YoloParam param, bool is_cuda);
	int detect(cv::Mat& image, std::vector<YoloDetection>& output);
    std::vector<std::string> mClassNames;

private:
	//加载模型
	int load_net(std::string str_path, cv::dnn::Net& net, bool is_cuda);

	//加载类别
	int load_class_list(std::string str_path, std::vector<std::string>&vec_class);

	cv::Mat format_yolov5(const cv::Mat& input);			//缩放为正方形图片

	std::string mNetModelPath;				//模型文件
    std::string mClassPath;					//类别文件
	bool mUse_cuda;

	cv::dnn::Net mNet;

	bool mInited = false;

    float INPUT_WIDTH = (float)640.0;				//网络width
    float INPUT_HEIGHT = (float)640.0;				//网络height
    float SCORE_THRESHOLD = (float)0.2;
    float NMS_THRESHOLD = (float)0.8;
    float CONFIDENCE_THRESHOLD = (float)0.4;
};
